import asyncio
from pyppeteer import launch
import os
from datetime import datetime
import logging
from typing import List
import json
import shutil
import sys
import tkinter as tk
from tkinter import ttk

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class TrackingCrawler:
    def __init__(self, output_dir: str = "tracking_screenshots"):
        self.output_dir = output_dir
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
    async def init_browser(self):
        """初始化浏览器"""
        self.browser = await launch(
            headless=True,  # 无头模式
            args=['--no-sandbox', '--disable-setuid-sandbox', '--disable-web-security']
        )
        
    async def close_browser(self):
        """关闭浏览器"""
        if hasattr(self, 'browser'):
            await self.browser.close()
            
    async def track_single(self, tracking_number: str) -> str:
        """使用SDK查询单个物流单号并截图"""
        try:
            # 创建一个新的页面
            page = await self.browser.newPage()
            await page.setViewport({'width': 1920, 'height': 1080})
            
            # 设置页面内容，包含SDK
            html_content = '''
                <!DOCTYPE html>
                <html>
                <head>
                    <meta charset="UTF-8">
                    <title>物流查询</title>
                    <script>
                        // 等待SDK加载完成的函数
                        function waitForSDK() {
                            return new Promise(function(resolve) {
                                function check() {
                                    if (window.YQV5) {
                                        resolve();
                                    } else {
                                        setTimeout(check, 100);
                                    }
                                }
                                check();
                            });
                        }
                    </script>
                    <script src="https://www.17track.net/externalcall.js"></script>
                </head>
                <body>
                    <div id="YQContainer"></div>
                    <script>
                        // 页面加载完成后执行查询
                        window.onload = function() {
                            waitForSDK().then(function() {
                                YQV5.trackSingle({
                                    YQ_ContainerId: "YQContainer",
                                    YQ_Height: 800,
                                    YQ_Fc: "0",
                                    YQ_Lang: "zh-cn",
                                    YQ_Num: "%s"
                                });
                            });
                        };
                    </script>
                </body>
                </html>
            ''' % tracking_number
            
            await page.setContent(html_content)
            # 捕获页面console.log输出
            def handle_console(msg):
                print('[浏览器console]', msg.text)
            page.on('console', handle_console)
            
            # 等待物流信息加载完成
            try:
                await page.waitForSelector('#YQContainer', {'timeout': 10000})
                # 额外等待以确保内容完全加载
                await asyncio.sleep(5)
                # 1. 找到包含 .ori-block 的 frame
                target_frame = None
                for frame in page.frames:
                    try:
                        has_ori = await frame.J('.ori-block')
                        if has_ori:
                            target_frame = frame
                            break
                    except Exception:
                        continue

                if target_frame:
                    # 2. 在该frame中获取高度，并设置iframe高度
                    height = await target_frame.evaluate('''() => {
                        const oriBlock = document.querySelector('.ori-block');
                        return oriBlock ? oriBlock.offsetHeight : null;
                    }''')
                    print('[浏览器console] oriBlock高度:', height)
                    if height:
                        # 3. 在主页面设置iframe高度
                        await page.evaluate(f'''
                            () => {{
                                const iframe = document.querySelector('iframe');
                                if (iframe) {{
                                    iframe.style.height = '{height + 500}px';
                                    iframe.removeAttribute('max-height');
                                }}
                            }}
                        ''')
                else:
                    print('[浏览器console] 未找到含.ori-block的frame')
            except Exception as e:
                logger.warning(f"等待选择器超时，但继续尝试截图: {str(e)}")
            
            # 生成文件名
            filename = f"{tracking_number}.png"

            # 获取#YQContainer的位置和大小，只截取有内容的区域
            clip = await page.evaluate('''() => {
                const rect = document.getElementById('YQContainer').getBoundingClientRect();
                return {
                    x: rect.x,
                    y: rect.y,
                    width: rect.width,
                    height: rect.height
                };
            }''')

            # 截图
            await page.screenshot({
                'path': filename,
                'fullPage': False,
                'clip': clip
            })

            # 遍历所有iframe，查找包含.text-uppercase的frame并抓取信息
            info = None
            for frame in page.frames:
                try:
                    has_danhao = await frame.J('.text-uppercase')
                    if has_danhao:
                        info = await frame.evaluate('''() => {
                            // 单号
                            let danhao = '';
                            let danhaoEl = document.querySelector('.text-uppercase');
                            if (danhaoEl && danhaoEl.innerText) {
                                danhao = danhaoEl.innerText.trim();
                            }
                            
                            // 状态（第一个.text-capitalize）
                            let zhuangtai = '';
                            let capList = Array.from(document.querySelectorAll('.text-capitalize'));
                            if (capList.length > 0 && capList[0].innerText) {
                                zhuangtai = capList[0].innerText.trim();
                            }
                            
                            // 如果没有找到状态，尝试从其他元素获取
                            if (!zhuangtai) {
                                let statusElements = document.querySelectorAll('.status-text, .status-label, .status');
                                for (let el of statusElements) {
                                    if (el && el.innerText && el.innerText.trim()) {
                                        zhuangtai = el.innerText.trim();
                                        break;
                                    }
                                }
                            }
                            
                            // 轨迹
                            let guiji = [];
                            let oriBlock = document.querySelector('.ori-block');
                            if (oriBlock) {
                                let ddList = oriBlock.querySelectorAll('dd');  // 抓取所有dd标签
                                ddList.forEach(dd => {
                                    let time = '';
                                    let content = '';
                                    
                                    let timeEl = dd.querySelector('time');
                                    if (timeEl && timeEl.innerText) {
                                        time = timeEl.innerText.trim();
                                    }
                                    
                                    let contentEl = dd.querySelector('p');
                                    if (contentEl && contentEl.innerText) {
                                        content = contentEl.innerText.trim();
                                    }
                                    
                                    if (time || content) {  // 只添加非空的记录
                                        guiji.push({
                                            '时间': time,
                                            '内容': content
                                        });
                                    }
                                });
                            }
                            
                            // 获取当前时间
                            let now = new Date();
                            let queryTime = now.getFullYear() + '-' + 
                                          String(now.getMonth() + 1).padStart(2, '0') + '-' + 
                                          String(now.getDate()).padStart(2, '0') + ' ' + 
                                          String(now.getHours()).padStart(2, '0') + ':' + 
                                          String(now.getMinutes()).padStart(2, '0') + ':' + 
                                          String(now.getSeconds()).padStart(2, '0');
                            
                            return {
                                '单号': danhao || '%s',  // 如果没找到单号，使用传入的单号
                                '状态': zhuangtai || '未知状态',
                                '轨迹': guiji,
                                '查询时间': queryTime
                            };
                        }''' % tracking_number)
                        break
                except Exception as e:
                    logger.error(f"提取信息时发生错误: {str(e)}")
                    continue
            if info is None:
                info = {
                    '单号': tracking_number,
                    '状态': '未知状态',
                    '轨迹': [],
                    '查询时间': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                }

            # 关闭页面
            await page.close()

            logger.info(f"成功查询并截图: {tracking_number}")
            return filename, info
            
        except Exception as e:
            logger.error(f"查询 {tracking_number} 时发生错误: {str(e)}")
            if 'page' in locals():
                await page.close()
            return None
            
    async def track_multiple(self, tracking_numbers: List[str]):
        """批量查询多个物流单号"""
        try:
            await self.init_browser()
            results = []
            info_list = []
            total = len(tracking_numbers)
            log_file = 'tracking_progress.log'
            with open(log_file, 'w', encoding='utf-8') as f:
                f.write('')
            # 查询所有已存在的图片名（不含扩展名）
            existing_numbers = set()
            output_dir = self.output_dir
            if os.path.exists(output_dir):
                for folder in os.listdir(output_dir):
                    folder_path = os.path.join(output_dir, folder)
                    if os.path.isdir(folder_path):
                        for f in os.listdir(folder_path):
                            if f.lower().endswith('.png'):
                                existing_numbers.add(os.path.splitext(f)[0])
            # 确保jsons文件夹存在
            jsons_dir = os.path.abspath('jsons')
            if not os.path.exists(jsons_dir):
                os.makedirs(jsons_dir)
            for idx, number in enumerate(tracking_numbers):
                if os.path.exists('tracking_stop.flag'):
                    logger.info("检测到暂停标志，主动退出")
                    break
                if str(number) in existing_numbers:
                    log_line = f"{idx+1}/{total} 跳过 {number}，已存在同名图片\n"
                    with open(log_file, 'a', encoding='utf-8') as f:
                        f.write(log_line)
                    continue
                result = await self.track_single(number)
                log_line = ''
                if result:
                    filename, info = result
                    results.append(filename)
                    info_list.append(info)
                    
                    # 处理状态文件夹名称
                    status = info.get('状态', '未知状态').strip()
                    # 标准化状态名称
                    status_map = {
                        '运输途中': ['运输途中', '运验途中', '物流途中', '配送途中', '运蛆途中', '运辆途', '运益途中', '途中'],
                        '收到信息': ['收到信息', '已收到信息', '信息已收到', 'notified'],
                        '拖动下方滑块完成拼图': ['拖动下方滑块完成拼图', '拖动滑块', '完成拼图'],
                        '查询不到': ['查询不到', '无法查询', '未查询到', '查间不到', '坦和不到', '请稍后查询', '于运输商网络'],
                        '成功签收': ['成功签收', '签收成功'],
                        '派送异常': ['派送异常', '配送异常', '投递异常'],
                        '查询中': ['查询中', '正在查询', '查询进行中', '查间中']
                    }
                    
                    # 标准化状态
                    normalized_status = '未知状态'
                    for std_status, variants in status_map.items():
                        if any(variant in status for variant in variants):
                            normalized_status = std_status
                            break
                    
                    # 使用标准化后的状态作为文件夹名
                    status_folder = normalized_status
                    subdir = os.path.join(self.output_dir, status_folder)
                    if not os.path.exists(subdir):
                        os.makedirs(subdir)
                    dest_path = os.path.join(subdir, os.path.basename(filename))
                    shutil.move(filename, dest_path)
                    
                    # 保存 JSON 文件，使用单号作为文件名
                    json_filename = f"{info.get('单号', number)}.json"
                    info_json_path = os.path.join(jsons_dir, json_filename)
                    with open(info_json_path, 'w', encoding='utf-8') as f:
                        json.dump(info, f, ensure_ascii=False, indent=2)
                    
                    log_line = f"{idx+1}/{total} 查询 {number} 成功，状态：{normalized_status}，文件路径：{os.path.abspath(dest_path)}\n"
                else:
                    log_line = f"{idx+1}/{total} 查询 {number} 失败\n"
                with open(log_file, 'a', encoding='utf-8') as f:
                    f.write(log_line)
                await asyncio.sleep(0.5)
            return results
        finally:
            await self.close_browser()
            # 退出时删除标志文件
            if os.path.exists('tracking_stop.flag'):
                os.remove('tracking_stop.flag')

async def main():
    # 支持命令行参数传入json文件
    tracking_numbers = []
    if len(sys.argv) > 1:
        json_file = sys.argv[1]
        try:
            with open(json_file, 'r', encoding='utf-8') as f:
                tracking_numbers = json.load(f)
            logger.info(f"从文件 {json_file} 读取到 {len(tracking_numbers)} 个单号")
        except Exception as e:
            logger.error(f"读取单号文件失败: {e}")
            return
    else:
        tracking_numbers = [
            '9200190381903610937180',
            '9200190373170070436640',
            '9200190373170070446434',
            '9200190381903610937181',
        ]
    
    crawler = TrackingCrawler()
    results = await crawler.track_multiple(tracking_numbers)
    
    logger.info(f"完成查询，共处理 {len(results)} 个单号")
    logger.info(f"截图保存在: {crawler.output_dir}")

if __name__ == "__main__":
    asyncio.run(main())

